import collections

from sortedcontainers import SortedList


class MKAverage:

    def __init__(self, m: int, k: int):
        self._m, self._k = m, k
        self._size = 0  # 当前元素数量
        self._queue = collections.deque()  # 数据流中的数值列表
        self._smaller = SortedList()  # 最小的k个值
        self._middle = SortedList()  # 中间用于计算的k个值
        self._bigger = SortedList()  # 最大的k个值
        self._sum = 0  # 中间部分容器的和

    def addElement(self, num1: int) -> None:
        # 当前数据流中的整数还少于m个的情况
        if self._size < self._m:
            self._queue.append(num1)
            self._size += 1
            if self._size == self._m:
                queue = sorted(self._queue)
                for i in range(self._k):
                    self._smaller.add(queue[i])
                for i in range(self._k, self._m - self._k):
                    self._middle.add(queue[i])
                for i in range(self._m - self._k, self._m):
                    self._bigger.add(queue[i])
            self._sum = sum(self._middle)

        # 当前数据流中的整数达到m个的情况
        else:
            self._queue.append(num1)
            num2 = self._queue.popleft()
            if num1 <= self._smaller[-1]:
                self._smaller.add(num1)
                if num2 in self._smaller:
                    self._smaller.remove(num2)
                elif num2 in self._bigger:
                    self._change_l_to_m()
                    self._change_m_to_r()
                    self._bigger.remove(num2)
                else:
                    self._change_l_to_m()
                    self._middle.remove(num2)
                    self._sum -= num2
            elif num1 >= self._bigger[-1]:
                self._bigger.add(num1)
                if num2 in self._smaller:
                    self._change_r_to_m()
                    self._change_m_to_l()
                    self._smaller.remove(num2)
                elif num2 in self._bigger:
                    self._bigger.remove(num2)
                else:
                    self._change_r_to_m()
                    self._middle.remove(num2)
                    self._sum -= num2
            else:
                self._middle.add(num1)
                self._sum += num1
                if num2 in self._smaller:
                    self._change_m_to_l()
                    self._smaller.remove(num2)
                elif num2 in self._bigger:
                    self._change_m_to_r()
                    self._bigger.remove(num2)
                else:
                    self._middle.remove(num2)
                    self._sum -= num2

    def _change_l_to_m(self):
        num = self._smaller.pop()
        self._middle.add(num)
        self._sum += num

    def _change_m_to_r(self):
        num = self._middle.pop()
        self._bigger.add(num)
        self._sum -= num

    def _change_m_to_l(self):
        num = self._middle.pop(0)
        self._smaller.add(num)
        self._sum -= num

    def _change_r_to_m(self):
        num = self._bigger.pop(0)
        self._middle.add(num)
        self._sum += num

    def calculateMKAverage(self) -> int:
        # 当前数据流中的整数还少于m个的情况
        if self._size < self._m:
            return -1

        # 当前数据流中的整数达到m个的情况
        else:
            return int(self._sum / (self._m - 2 * self._k))


if __name__ == "__main__":
    obj = MKAverage(3, 1)
    obj.addElement(3)
    obj.addElement(1)
    print(obj.calculateMKAverage())  # -1
    obj.addElement(10)
    print(obj.calculateMKAverage())  # 3
    obj.addElement(5)
    obj.addElement(5)
    obj.addElement(5)
    print(obj.calculateMKAverage())  # 5
